from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import os
import os.path

import re
import sys
import tarfile
from six.moves import urllib

import json

import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt
from matplotlib.font_manager import FontProperties
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont

from utils import visualization_utils as vis_util
from utils import label_map_util

def load_image_into_numpy_array(image):
    (im_width, im_height) = image.size
    return np.array(image.getdata()).reshape(
        (im_height, im_width, 3)).astype(np.uint8), im_height, im_width

def vehicle_detection(test_img_path, NUM_CLASSES, PATH_TO_CKPT, PATH_TO_LABELS, image_detection=None):
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')

    label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
    categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
    category_index = label_map_util.create_category_index(categories)
    with detection_graph.as_default():
        with tf.Session(graph=detection_graph) as sess:
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')
            image = Image.open(test_img_path)
            #image_n = load_image_into_numpy_array(image)
            image_np = load_image_into_numpy_array(image)
            image_np = image_np[0]
            image_np_expanded = np.expand_dims(image_np, axis=0)
            (boxes, scores, classes, num) = sess.run(
                [detection_boxes, detection_scores, detection_classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
            vis_util.visualize_boxes_and_labels_on_image_array(
                image_np,
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                category_index,
                use_normalized_coordinates=True,
                line_thickness=8)
            image_detecion = Image.fromarray(image_np,'RGB')
            #image_detecion.save(image_detection)
            return scores,classes,boxes
